import tensorflow as tf
from tensorflow.keras import layers, models
import numpy as np
import matplotlib
# 使用非交互式后端避免PyCharm兼容性问题
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import matplotlib.font_manager
import os
import datetime
import time

# 禁用XLA优化以避免MPS兼容性问题
tf.config.optimizer.set_jit(False)

# 强制使用float32精度以提高MPS兼容性
tf.keras.mixed_precision.set_global_policy('float32')
print("使用默认精度(float32)")

# 初始化TensorBoard可用性标志
TENSORBOARD_AVAILABLE = False

# 尝试导入TensorBoard，如果失败则跳过
try:
    from tensorflow.keras.callbacks import TensorBoard
    TENSORBOARD_AVAILABLE = True
except ImportError:
    print("TensorBoard不可用，将跳过相关功能")

# MPS优化配置
def configure_mps():
    """配置MPS优化设置"""
    gpus = tf.config.list_physical_devices('GPU')
    if gpus:
        try:
            # 设置GPU内存增长
            for gpu in gpus:
                tf.config.experimental.set_memory_growth(gpu, True)
            print(f"检测到 {len(gpus)} 个GPU设备，已设置内存增长模式")
            
            # MPS特定优化
            # 启用内存增长和优化
            for gpu in gpus:
                tf.config.experimental.set_memory_growth(gpu, True)
                # 设置虚拟设备配置以更好地管理内存
                # tf.config.set_logical_device_configuration(
                #     gpu,
                #     [tf.config.LogicalDeviceConfiguration(memory_limit=4096)])
            
            return True
        except RuntimeError as e:
            print(f"设置GPU配置失败: {e}")
            return False
    else:
        print("未检测到GPU设备，将使用CPU进行计算")
        return False

# 设置字体以正确显示负号
plt.rcParams['axes.unicode_minus'] = False

def load_local_mnist(data_path="./mnist_data"):
    """
    从本地目录加载MNIST数据
    """
    print(f"尝试从本地目录 {data_path} 加载MNIST数据...")
    
    # 检查是否存在本地MNIST数据文件
    mnist_npz_path = os.path.join(data_path, "mnist.npz") if os.path.isdir(data_path) else data_path
    
    if os.path.exists(mnist_npz_path) and mnist_npz_path.endswith('.npz'):
        try:
            # 从本地NPZ文件加载数据
            with np.load(mnist_npz_path, allow_pickle=True) as f:
                x_train, y_train = f['x_train'], f['y_train']
                x_test, y_test = f['x_test'], f['y_test']
            print(f"成功从本地文件 {mnist_npz_path} 加载数据")
            return (x_train, y_train), (x_test, y_test)
        except Exception as e:
            print(f"从本地NPZ文件加载数据失败: {e}")
    
    # 如果本地加载失败，使用tf.keras方式加载
    print("使用TensorFlow内置方式加载MNIST数据...")
    return tf.keras.datasets.mnist.load_data()

def load_and_preprocess_data():
    """加载并预处理MNIST数据"""
    print("正在加载MNIST数据集...")
    (x_train, y_train), (x_test, y_test) = load_local_mnist("./mnist_data")
    
    # 数据预处理 - 使用向量化操作提高效率
    x_train = np.expand_dims(x_train.astype('float32') / 255.0, axis=-1)
    x_test = np.expand_dims(x_test.astype('float32') / 255.0, axis=-1)
    
    # 标签one-hot编码
    y_train = tf.keras.utils.to_categorical(y_train, 10).astype('float32')
    y_test = tf.keras.utils.to_categorical(y_test, 10).astype('float32')
    
    print(f"训练集形状: {x_train.shape}")
    print(f"测试集形状: {x_test.shape}")
    print(f"训练标签形状: {y_train.shape}")
    print(f"测试标签形状: {y_test.shape}")
    
    return (x_train, y_train), (x_test, y_test)

def create_cnn_model():
    """创建优化的卷积神经网络模型"""
    model = models.Sequential([
        # 第一个卷积层 - 降低滤波器数量以提高MPS兼容性
        layers.Conv2D(32, (3, 3), activation='relu', input_shape=(28, 28, 1)),
        layers.BatchNormalization(),  # 添加批归一化以加速训练
        layers.MaxPooling2D((2, 2)),
        
        # 第二个卷积层
        layers.Conv2D(64, (3, 3), activation='relu'),
        layers.BatchNormalization(),
        layers.MaxPooling2D((2, 2)),
        
        # 第三个卷积层
        layers.Conv2D(64, (3, 3), activation='relu'),
        layers.BatchNormalization(),
        
        # 展平层
        layers.Flatten(),
        
        # 全连接层 - 适度增加神经元数量
        layers.Dense(128, activation='relu'),
        layers.Dropout(0.4),  # 适度dropout以防止过拟合
        
        # 输出层 - 明确指定float32以提高MPS兼容性
        layers.Dense(10, activation='softmax', dtype='float32')
    ])
    
    return model

def visualize_samples(x_data, y_data, num_samples=9):
    """可视化数据样本"""
    plt.figure(figsize=(8, 8))
    for i in range(num_samples):
        plt.subplot(3, 3, i+1)
        plt.imshow(x_data[i].reshape(28, 28), cmap='gray')
        plt.title(f'标签: {np.argmax(y_data[i])}')
        plt.axis('off')
    plt.tight_layout()
    # 修复matplotlib tostring_rgb错误
    try:
        plt.savefig('mnist_samples_mps.png')
    except AttributeError:
        print("保存图像时遇到matplotlib兼容性问题，跳过保存")
    plt.show()
    plt.close()  # 关闭图形以释放内存

def plot_training_history(history):
    """绘制训练历史"""
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4))
    
    # 绘制损失曲线
    ax1.plot(history.history['loss'], label='训练损失')
    ax1.plot(history.history['val_loss'], label='验证损失')
    ax1.set_title('模型损失')
    ax1.set_xlabel('轮次')
    ax1.set_ylabel('损失')
    ax1.legend()
    ax1.grid(True)
    
    # 绘制准确率曲线
    ax2.plot(history.history['accuracy'], label='训练准确率')
    ax2.plot(history.history['val_accuracy'], label='验证准确率')
    ax2.set_title('模型准确率')
    ax2.set_xlabel('轮次')
    ax2.set_ylabel('准确率')
    ax2.legend()
    ax2.grid(True)
    
    plt.tight_layout()
    # 修复matplotlib tostring_rgb错误
    try:
        plt.savefig('training_history_mps.png')
    except AttributeError:
        print("保存图像时遇到matplotlib兼容性问题，跳过保存")
    plt.show()
    plt.close()  # 关闭图形以释放内存

def benchmark_device_performance():
    """基准测试设备性能"""
    print("进行设备性能基准测试...")
    
    # 创建测试张量
    size = 5000
    with tf.device('/CPU:0'):
        a_cpu = tf.random.normal([size, size], dtype=tf.float32)
        b_cpu = tf.random.normal([size, size], dtype=tf.float32)
        
        # 预热
        _ = tf.matmul(a_cpu, b_cpu)
        
        start_time = time.time()
        c_cpu = tf.matmul(a_cpu, b_cpu)
        _ = tf.reduce_sum(c_cpu)
        cpu_time = time.time() - start_time
    
    print(f"CPU矩阵乘法时间 ({size}x{size}): {cpu_time:.4f} 秒")
    
    # 如果GPU可用，测试GPU性能
    gpus = tf.config.list_physical_devices('GPU')
    if gpus:
        with tf.device('/GPU:0'):
            a_gpu = tf.random.normal([size, size], dtype=tf.float32)
            b_gpu = tf.random.normal([size, size], dtype=tf.float32)
            
            # 预热
            _ = tf.matmul(a_gpu, b_gpu)
            
            start_time = time.time()
            c_gpu = tf.matmul(a_gpu, b_gpu)
            _ = tf.reduce_sum(c_gpu)
            gpu_time = time.time() - start_time
        
        print(f"GPU矩阵乘法时间 ({size}x{size}): {gpu_time:.4f} 秒")
        print(f"加速比: {cpu_time/gpu_time:.2f}x")
        return cpu_time, gpu_time
    else:
        print("GPU不可用，跳过GPU性能测试")
        return cpu_time, None

def main():
    """主函数"""
    print("MNIST手写数字识别 - MPS优化版")
    print("=" * 50)
    
    # 配置MPS优化
    mps_available = configure_mps()
    
    # 基准测试
    benchmark_device_performance()
    
    # 加载数据
    (x_train, y_train), (x_test, y_test) = load_and_preprocess_data()
    
    # 可视化样本数据
    print("\n可视化样本数据...")
    visualize_samples(x_train, y_train)
    
    # 创建模型
    print("\n创建卷积神经网络模型...")
    model = create_cnn_model()
    
    # 显示模型结构
    model.summary()
    
    # 编译模型 - 使用优化的优化器和学习率
    try:
        model.compile(
            optimizer=tf.keras.optimizers.Adam(learning_rate=0.001),
            loss='categorical_crossentropy',
            metrics=['accuracy']
        )
    except Exception as e:
        print(f"模型编译失败，使用默认设置: {e}")
        model.compile(
            optimizer='adam',
            loss='categorical_crossentropy',
            metrics=['accuracy']
        )

    # 设置TensorBoard回调（如果可用）
    callbacks = []
    tensorboard_enabled = TENSORBOARD_AVAILABLE
    log_dir = None
    
    if tensorboard_enabled:
        try:
            # 使用简单的英文路径，避免中文字符问题
            base_log_dir = "logs_fit_mps"
            # 确保基础日志目录存在
            os.makedirs(base_log_dir, exist_ok=True)
            # 创建基于时间戳的子目录
            timestamp = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
            log_dir = os.path.join(base_log_dir, timestamp)
            # 创建完整的日志目录
            os.makedirs(log_dir, exist_ok=True)
            tensorboard_callback = TensorBoard(log_dir=log_dir, histogram_freq=1)
            callbacks.append(tensorboard_callback)
            print(f"TensorBoard功能已启用，日志将保存到: {log_dir}")
        except Exception as e:
            print(f"创建TensorBoard日志目录失败: {e}")
            print("TensorBoard功能将被禁用")
            tensorboard_enabled = False
    
    # 添加早停回调以防止过拟合并节省时间
    early_stopping = tf.keras.callbacks.EarlyStopping(
        monitor='val_accuracy',
        patience=3,
        restore_best_weights=True
    )
    callbacks.append(early_stopping)
    
    # 训练模型
    print("\n开始训练模型...")
    start_time = time.time()
    
    # 明确指定在GPU上运行（如果可用）
    device = '/GPU:0' if mps_available else '/CPU:0'
    print(f"使用设备: {device}")
    
    with tf.device(device):
        history = model.fit(
            x_train, y_train,
            batch_size=128,  # 恢复到较小的批次大小以提高稳定性
            epochs=5,        # 减少训练轮数以避免长时间运行
            verbose=1,
            validation_data=(x_test, y_test),
            callbacks=callbacks
        )
    
    training_time = time.time() - start_time
    print(f"训练耗时: {training_time:.2f} 秒")
    
    # 绘制训练历史
    print("\n绘制训练历史...")
    plot_training_history(history)
    
    # 评估模型
    print("\n评估模型...")
    start_time = time.time()
    
    with tf.device(device):
        test_loss, test_accuracy = model.evaluate(x_test, y_test, verbose=0)
    
    evaluation_time = time.time() - start_time
    print(f"评估耗时: {evaluation_time:.4f} 秒")
    print(f"测试损失: {test_loss:.4f}")
    print(f"测试准确率: {test_accuracy:.4f}")
    
    # 预测示例
    print("\n预测示例...")
    with tf.device(device):
        predictions = model.predict(x_test[:10])
    
    predicted_classes = np.argmax(predictions, axis=1)
    true_classes = np.argmax(y_test[:10], axis=1)
    
    print("前10个测试样本的预测结果:")
    for i in range(10):
        print(f"样本 {i+1}: 真实标签={true_classes[i]}, 预测标签={predicted_classes[i]}, "
              f"置信度={np.max(predictions[i]):.4f}")
    
    # 保存模型
    model.save('cnn_mnist_model_mps.h5')
    print("\n模型已保存为 'cnn_mnist_model_mps.h5'")
    
    # 提供TensorBoard启动命令
    if tensorboard_enabled and callbacks:
        print("\n训练完成！要查看TensorBoard，请在终端中运行以下命令：")
        print(f"tensorboard --logdir {log_dir}")

if __name__ == "__main__":
    main()